# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: inkyu

from collections import namedtuple, deque

Entry = namedtuple('Entry', ['timestamp', 'qty', 'volume'])

QtyVol = namedtuple('QtyVol', ['qty', 'volume'])

Result = namedtuple('Result', ['price', 'qty'])


class VwapCalculator(object):
  def __init__(self, lifetime_sec):
    assert lifetime_sec in range(1, 1000), lifetime_sec
    self._lifetime = lifetime_sec * 1e9
    self._data = deque()

    self._total_qty = 0
    self._total_volume = 0

  def get(self):
    if len(self._data) > 0:
      duration = self._data[-1].timestamp - self._data[0].timestamp
      if duration > 0. and self._total_qty > 0.:
        vwap_price = self._total_volume / self._total_qty
        qty = self._total_qty / duration
        return Result(price=vwap_price, qty=qty)
    return None

  def set(self, *, volume, qty, timestamp):
    if len(self._data) > 0:
      back = self._data[-1]
      self._total_qty += back.qty
      self._total_volume += back.volume
      self._data.append(Entry(timestamp, qty, volume))

      initial_ts = timestamp - self._lifetime
      initial_value = None

      while len(self._data) > 0:
        current = self._data[0]
        if current.timestamp >= initial_ts:
          if current.timestamp == initial_ts:
            initial_value = None
          break
        else:
          initial_value = QtyVol(qty=current.qty, volume=current.volume)
          front = self._data.popleft()
          self._total_qty -= front.qty
          self._total_volume -= front.volume
          if self._total_qty <= 0. or self._total_volume <= 0.:
            self._total_qty = 0.
            self._total_volume = 0.
      if initial_value is not None:
        self._total_qty += initial_value.qty
        self._total_volume += initial_value.volume

        self._data.appendleft(
            Entry(timestamp=initial_ts, qty=initial_value.qty, volume=initial_value.volume))
    else:
      self._data.append(Entry(timestamp=timestamp, qty=qty, volume=volume))
